#pragma once
#include "SkrRenderGraph/frontend/base_types.hpp"
#include "SkrRenderGraph/frontend/blackboard.hpp"
#include "SkrContainersDef/hashmap.hpp"

#ifndef RG_MAX_FRAME_IN_FLIGHT
#define RG_MAX_FRAME_IN_FLIGHT 3
#endif

namespace skr
{
namespace render_graph
{
class SKR_RENDER_GRAPH_API RenderGraphProfiler
{
public:
    virtual ~RenderGraphProfiler() = default;
    virtual void on_acquire_executor(class RenderGraph&, class RenderGraphFrameExecutor&) {}
    virtual void on_cmd_begin(class RenderGraph&, class RenderGraphFrameExecutor&) {}
    virtual void on_cmd_end(class RenderGraph&, class RenderGraphFrameExecutor&) {}
    virtual void on_pass_begin(class RenderGraph&, class RenderGraphFrameExecutor&, class PassNode& pass) {}
    virtual void on_pass_end(class RenderGraph&, class RenderGraphFrameExecutor&, class PassNode& pass) {}
    virtual void before_commit(class RenderGraph&, class RenderGraphFrameExecutor&) {}
    virtual void after_commit(class RenderGraph&, class RenderGraphFrameExecutor&) {}
};

struct SKR_RENDER_GRAPH_API IRenderGraphPhase
{
    virtual ~IRenderGraphPhase() SKR_NOEXCEPT;
    virtual void on_execute(RenderGraph* graph, RenderGraphFrameExecutor* executor, RenderGraphProfiler* profiler) SKR_NOEXCEPT;

    skr::Vector<ResourceNode*>& get_resources(RenderGraph* graph) SKR_NOEXCEPT;
    skr::Vector<PassNode*>& get_passes(RenderGraph* graph) SKR_NOEXCEPT;
};

class SKR_RENDER_GRAPH_API RenderGraph
{
public:
    friend class RenderGraphViz;
    class SKR_RENDER_GRAPH_API RenderGraphBuilder
    {
    public:
        friend class RenderGraph;
        friend class RenderGraphBackend;
        RenderGraphBuilder& frontend_only() SKR_NOEXCEPT;
        RenderGraphBuilder& backend_api(ECGPUBackend backend) SKR_NOEXCEPT;
        RenderGraphBuilder& with_device(CGPUDeviceId device) SKR_NOEXCEPT;
        RenderGraphBuilder& with_gfx_queue(CGPUQueueId queue) SKR_NOEXCEPT;
        RenderGraphBuilder& with_cmpt_queues(const skr::Vector<CGPUQueueId>& queues) SKR_NOEXCEPT;
        RenderGraphBuilder& with_cpy_queues(const skr::Vector<CGPUQueueId>& queues) SKR_NOEXCEPT;
        RenderGraphBuilder& enable_memory_aliasing() SKR_NOEXCEPT;

    protected:
        bool memory_aliasing = false;
        bool no_backend = false;
        ECGPUBackend api;
        CGPUDeviceId device;
        CGPUQueueId gfx_queue;
        skr::Vector<CGPUQueueId> cmpt_queues;
        skr::Vector<CGPUQueueId> cpy_queues;
    };
    using RenderGraphSetupFunction = skr::stl_function<void(class RenderGraph::RenderGraphBuilder&)>;
    static RenderGraph* create(const RenderGraphSetupFunction& setup) SKR_NOEXCEPT;
    static void destroy(RenderGraph* g) SKR_NOEXCEPT;
    
    class SKR_RENDER_GRAPH_API RenderPassBuilder
    {
    public:
        friend class RenderGraph;
        RenderPassBuilder& set_name(const char8_t* name) SKR_NOEXCEPT;
        // textures
        RenderPassBuilder& read(const char8_t* name, TextureSRVHandle handle) SKR_NOEXCEPT;
        RenderPassBuilder& write(uint32_t mrt_index, TextureRTVHandle handle,
            ECGPULoadAction load_action = CGPU_LOAD_ACTION_CLEAR,
            CGPUClearValue clear_color = fastclear_0000,
            ECGPUStoreAction store_action = CGPU_STORE_ACTION_STORE) SKR_NOEXCEPT;
        RenderPassBuilder& resolve_msaa(uint32_t mrt_index, TextureSubresourceHandle handle);
        
        RenderPassBuilder& set_depth_stencil(TextureDSVHandle handle,
            ECGPULoadAction dload_action = CGPU_LOAD_ACTION_CLEAR,
            ECGPUStoreAction dstore_action = CGPU_STORE_ACTION_STORE,
            ECGPULoadAction sload_action = CGPU_LOAD_ACTION_CLEAR,
            ECGPUStoreAction sstore_action = CGPU_STORE_ACTION_STORE) SKR_NOEXCEPT;

        // buffers
        RenderPassBuilder& read(const char8_t* name, BufferRangeHandle handle) SKR_NOEXCEPT;
        RenderPassBuilder& write(const char8_t* name, BufferRangeHandle handle) SKR_NOEXCEPT;
        RenderPassBuilder& use_buffer(PipelineBufferHandle buffer, ECGPUResourceState requested_state) SKR_NOEXCEPT;

        RenderPassBuilder& set_pipeline(CGPURenderPipelineId pipeline) SKR_NOEXCEPT;
        RenderPassBuilder& set_root_signature(CGPURootSignatureId signature) SKR_NOEXCEPT;
        
        // Performance hints
        RenderPassBuilder& with_flags(EPassFlags flags) SKR_NOEXCEPT;
    protected:
        RenderPassBuilder(RenderGraph& graph, RenderPassNode& pass) SKR_NOEXCEPT;
        RenderGraph& graph;
        RenderPassNode& node;
    };
    using RenderPassSetupFunction = skr::stl_function<void(RenderGraph&, class RenderGraph::RenderPassBuilder&)>;
    PassHandle add_render_pass(const RenderPassSetupFunction& setup, const RenderPassExecuteFunction& executor) SKR_NOEXCEPT;

    class SKR_RENDER_GRAPH_API ComputePassBuilder
    {
    public:
        friend class RenderGraph;
        ComputePassBuilder& set_name(const char8_t* name) SKR_NOEXCEPT;
        ComputePassBuilder& read(const char8_t* name, TextureSRVHandle handle) SKR_NOEXCEPT;

        ComputePassBuilder& readwrite(const char8_t* name, TextureUAVHandle handle) SKR_NOEXCEPT;
        
        ComputePassBuilder& read(const char8_t* name, BufferRangeHandle handle) SKR_NOEXCEPT;
        ComputePassBuilder& read(const char8_t* name, BufferHandle handle) SKR_NOEXCEPT;

        ComputePassBuilder& readwrite(const char8_t* name, BufferRangeHandle handle) SKR_NOEXCEPT;
        ComputePassBuilder& readwrite(const char8_t* name, BufferHandle handle) SKR_NOEXCEPT;

        ComputePassBuilder& read(const char8_t* name, AccelerationStructureSRVHandle handle) SKR_NOEXCEPT;

        ComputePassBuilder& set_pipeline(CGPUComputePipelineId pipeline) SKR_NOEXCEPT;
        ComputePassBuilder& set_root_signature(CGPURootSignatureId signature) SKR_NOEXCEPT;
        
        // Performance hints
        ComputePassBuilder& with_flags(EPassFlags flags) SKR_NOEXCEPT;

    protected:
        ComputePassBuilder(RenderGraph& graph, ComputePassNode& pass) SKR_NOEXCEPT;
        RenderGraph& graph;
        ComputePassNode& node;
    };
    using ComputePassSetupFunction = skr::stl_function<void(RenderGraph&, class RenderGraph::ComputePassBuilder&)>;
    PassHandle add_compute_pass(const ComputePassSetupFunction& setup, const ComputePassExecuteFunction& executor) SKR_NOEXCEPT;

    class SKR_RENDER_GRAPH_API CopyPassBuilder
    {
    public:
        friend class RenderGraph;
        CopyPassBuilder& set_name(const char8_t* name) SKR_NOEXCEPT;
        CopyPassBuilder& can_be_lone() SKR_NOEXCEPT;
        CopyPassBuilder& texture_to_texture(TextureSubresourceHandle src, TextureSubresourceHandle dst, ECGPUResourceState out_state = CGPU_RESOURCE_STATE_COPY_DEST) SKR_NOEXCEPT;
        CopyPassBuilder& buffer_to_buffer(BufferRangeHandle src, BufferRangeHandle dst, ECGPUResourceState out_state = CGPU_RESOURCE_STATE_COPY_DEST) SKR_NOEXCEPT;
        CopyPassBuilder& buffer_to_texture(BufferRangeHandle src, TextureSubresourceHandle dst, ECGPUResourceState out_state = CGPU_RESOURCE_STATE_COPY_DEST) SKR_NOEXCEPT;
        CopyPassBuilder& from_buffer(BufferRangeHandle src) SKR_NOEXCEPT;
        
        // Performance hints
        CopyPassBuilder& with_flags(EPassFlags flags) SKR_NOEXCEPT;

    protected:
        CopyPassBuilder(RenderGraph& graph, CopyPassNode& pass) noexcept;
        RenderGraph& graph;
        CopyPassNode& node;
    };
    using CopyPassSetupFunction = skr::stl_function<void(RenderGraph&, class RenderGraph::CopyPassBuilder&)>;
    PassHandle add_copy_pass(const CopyPassSetupFunction& setup, const CopyPassExecuteFunction& executor) SKR_NOEXCEPT;

    class SKR_RENDER_GRAPH_API PresentPassBuilder
    {
    public:
        friend class RenderGraph;

        PresentPassBuilder& set_name(const char8_t* name) SKR_NOEXCEPT;
        PresentPassBuilder& swapchain(CGPUSwapChainId chain, uint32_t index) SKR_NOEXCEPT;
        PresentPassBuilder& texture(TextureHandle texture, bool is_backbuffer = true) SKR_NOEXCEPT;

    protected:
        PresentPassBuilder(RenderGraph& graph, PresentPassNode& present) noexcept;
        RenderGraph& graph;
        PresentPassNode& node;
    };
    using PresentPassSetupFunction = skr::stl_function<void(RenderGraph&, class RenderGraph::PresentPassBuilder&)>;
    PassHandle add_present_pass(const PresentPassSetupFunction& setup) SKR_NOEXCEPT;

    class SKR_RENDER_GRAPH_API BufferBuilder
    {
    public:
        friend class RenderGraph;
        BufferBuilder& set_name(const char8_t* name) SKR_NOEXCEPT;
        BufferBuilder& with_tags(uint32_t tags) SKR_NOEXCEPT;
        BufferBuilder& import(CGPUBufferId buffer, ECGPUResourceState init_state) SKR_NOEXCEPT;
        BufferBuilder& allocate_dedicated() SKR_NOEXCEPT;
        BufferBuilder& structured(uint64_t first_element, uint64_t element_count, uint64_t element_stride) SKR_NOEXCEPT;
        BufferBuilder& size(uint64_t size) SKR_NOEXCEPT;
        BufferBuilder& with_flags(CGPUBufferFlags flags) SKR_NOEXCEPT;
        BufferBuilder& memory_usage(ECGPUMemoryUsage mem_usage) SKR_NOEXCEPT;
        BufferBuilder& allow_shader_readwrite() SKR_NOEXCEPT;
        BufferBuilder& allow_shader_read() SKR_NOEXCEPT;
        BufferBuilder& as_upload_buffer() SKR_NOEXCEPT;
        BufferBuilder& as_vertex_buffer() SKR_NOEXCEPT;
        BufferBuilder& as_index_buffer() SKR_NOEXCEPT;
        BufferBuilder& as_uniform_buffer() SKR_NOEXCEPT;
        BufferBuilder& prefer_on_device() SKR_NOEXCEPT;
        BufferBuilder& prefer_on_host() SKR_NOEXCEPT;

    protected:
        BufferBuilder(RenderGraph& graph, BufferNode& node) SKR_NOEXCEPT;
        RenderGraph& graph;
        BufferNode& node;
    };
    using BufferSetupFunction = skr::stl_function<void(RenderGraph&, class RenderGraph::BufferBuilder&)>;
    BufferHandle create_buffer(const BufferSetupFunction& setup) SKR_NOEXCEPT;
    BufferHandle get_buffer(const char8_t* name) SKR_NOEXCEPT;
    BufferHandle get_imported(CGPUBufferId buffer) SKR_NOEXCEPT;
    const ECGPUResourceState get_lastest_state(const BufferNode* buffer, const PassNode* pending_pass) const SKR_NOEXCEPT;

    class SKR_RENDER_GRAPH_API TextureBuilder
    {
    public:
        friend class RenderGraph;
        TextureBuilder& set_name(const char8_t* name) SKR_NOEXCEPT;
        TextureBuilder& with_flags(CGPUTextureFlags tags) SKR_NOEXCEPT;
        TextureBuilder& with_tags(uint32_t tags) SKR_NOEXCEPT;
        TextureBuilder& import(CGPUTextureId texture, ECGPUResourceState init_state) SKR_NOEXCEPT;
        TextureBuilder& extent(uint64_t width, uint64_t height, uint64_t depth = 1) SKR_NOEXCEPT;
        TextureBuilder& format(ECGPUFormat format) SKR_NOEXCEPT;
        TextureBuilder& array(uint32_t size) SKR_NOEXCEPT;
        TextureBuilder& sample_count(ECGPUSampleCount count) SKR_NOEXCEPT;
        TextureBuilder& allow_render_target() SKR_NOEXCEPT;
        TextureBuilder& allow_depth_stencil() SKR_NOEXCEPT;
        TextureBuilder& allow_readwrite() SKR_NOEXCEPT;
        TextureBuilder& allocate_dedicated() SKR_NOEXCEPT;
        TextureBuilder& allow_lone() SKR_NOEXCEPT;

    protected:
        TextureBuilder(RenderGraph& graph, TextureNode& node) SKR_NOEXCEPT;
        RenderGraph& graph;
        TextureNode& node;
        CGPUTextureId imported = nullptr;
    };
    using TextureSetupFunction = skr::stl_function<void(RenderGraph&, class RenderGraph::TextureBuilder&)>;
    TextureHandle create_texture(const TextureSetupFunction& setup) SKR_NOEXCEPT;
    TextureHandle get_texture(const char8_t* name) SKR_NOEXCEPT;
    TextureHandle get_imported(CGPUTextureId texture) SKR_NOEXCEPT;
    const ECGPUResourceState get_lastest_state(const TextureNode* texture, const PassNode* pending_pass) const SKR_NOEXCEPT;

    class SKR_RENDER_GRAPH_API AccelerationStructureBuilder
    {
    public:
        friend class RenderGraph;
        AccelerationStructureBuilder& set_name(const char8_t* name) SKR_NOEXCEPT;
        AccelerationStructureBuilder& import(CGPUAccelerationStructureId acceleration_structure) SKR_NOEXCEPT;

    protected:
        AccelerationStructureBuilder(RenderGraph& graph, AccelerationStructureNode& node) SKR_NOEXCEPT;
        RenderGraph& graph;
        AccelerationStructureNode& node;
        CGPUAccelerationStructureId imported = nullptr;
    };
    using AccelerationStructureSetupFunction = skr::stl_function<void(RenderGraph&, class RenderGraph::AccelerationStructureBuilder&)>;
    AccelerationStructureHandle create_acceleration_structure(const AccelerationStructureSetupFunction& setup) SKR_NOEXCEPT;
    AccelerationStructureHandle get_acceleration_structure(const char8_t* name) SKR_NOEXCEPT;
    AccelerationStructureHandle get_imported(CGPUAccelerationStructureId acceleration_structure) SKR_NOEXCEPT;
    const ECGPUResourceState get_lastest_state(const AccelerationStructureNode* acceleration_structure, const PassNode* pending_pass) const SKR_NOEXCEPT;

    BufferNode* resolve(BufferHandle hdl) SKR_NOEXCEPT; 
    TextureNode* resolve(TextureHandle hdl) SKR_NOEXCEPT;
    AccelerationStructureNode* resolve(AccelerationStructureHandle hdl) SKR_NOEXCEPT;
    PassNode* resolve(PassHandle hdl) SKR_NOEXCEPT;
    const CGPUBufferDescriptor* resolve_descriptor(BufferHandle hdl) SKR_NOEXCEPT;
    const CGPUTextureDescriptor* resolve_descriptor(TextureHandle hdl) SKR_NOEXCEPT;
    
    inline Blackboard& get_blackboard() SKR_NOEXCEPT
    {
        return *blackboard;
    }
    virtual CGPUDeviceId get_backend_device() SKR_NOEXCEPT { return nullptr; }
    virtual CGPUQueueId get_gfx_queue() SKR_NOEXCEPT { return nullptr; }
    virtual const skr::Vector<CGPUQueueId>& get_cmpt_queues() SKR_NOEXCEPT { static skr::Vector<CGPUQueueId> empty; return empty; }
    virtual const skr::Vector<CGPUQueueId>& get_cpy_queues() SKR_NOEXCEPT { static skr::Vector<CGPUQueueId> empty; return empty; }
    inline uint64_t get_frame_index() const SKR_NOEXCEPT { return frame_index; }
    inline struct NodeAndEdgeFactory* get_node_factory() SKR_NOEXCEPT { return node_factory; }
    virtual uint32_t collect_garbage(uint64_t critical_frame,
        uint32_t tex_with_tags = kRenderGraphDefaultResourceTag | kRenderGraphDynamicResourceTag, uint32_t tex_without_flags = 0,
        uint32_t buf_with_tags = kRenderGraphDefaultResourceTag | kRenderGraphDynamicResourceTag, uint32_t buf_without_flags = 0) SKR_NOEXCEPT
    {
        return collect_texture_garbage(critical_frame, tex_with_tags,tex_without_flags)
            + collect_buffer_garbage(critical_frame, buf_with_tags, buf_without_flags);
    }

    RenderGraph(const RenderGraphBuilder& builder) SKR_NOEXCEPT;
    virtual ~RenderGraph() SKR_NOEXCEPT = default;

    // interfaces
    friend struct IRenderGraphPhase;
    using BeforeExecuteCallback = skr::stl_function<void(RenderGraph&)>;
    void add_before_execute_callback(const BeforeExecuteCallback& callback);
    virtual void wait_frame(uint64_t frame_index) SKR_NOEXCEPT;
    virtual uint64_t execute(RenderGraphProfiler* profiler = nullptr) SKR_NOEXCEPT;
    virtual uint32_t collect_texture_garbage(uint64_t critical_frame,
        uint32_t with_tags = kRenderGraphDefaultResourceTag | kRenderGraphDynamicResourceTag, uint32_t without_flags = 0) SKR_NOEXCEPT { return 0; }
    virtual uint32_t collect_buffer_garbage(uint64_t critical_frame,
        uint32_t with_tags = kRenderGraphDefaultResourceTag | kRenderGraphDynamicResourceTag, uint32_t without_flags = 0) SKR_NOEXCEPT { return 0; }

    uint32_t foreach_textures(skr::stl_function<void(TextureNode*)> texture) SKR_NOEXCEPT;
    uint32_t foreach_passes(TextureHandle texture,
        skr::stl_function<void(PassNode* writer, TextureNode* tex, RenderGraphEdge* edge)>) const SKR_NOEXCEPT;
    uint32_t foreach_passes(BufferHandle buffer,
        skr::stl_function<void(PassNode* reader, BufferNode* buf, RenderGraphEdge* edge)>) const SKR_NOEXCEPT;
    const skr::Vector<PassNode*>& get_passes() const SKR_NOEXCEPT { return passes; }
    const skr::Vector<ResourceNode*>& get_resources() const SKR_NOEXCEPT { return resources; }

protected:
    virtual void initialize() SKR_NOEXCEPT;
    virtual void finalize() SKR_NOEXCEPT;

    bool aliasing_enabled = false;
    uint64_t frame_index = 0;

    struct NodeAndEdgeFactory* node_factory = nullptr;
    Blackboard* blackboard = nullptr;
    DependencyGraph* graph = nullptr;

    skr::Vector<PassNode*> passes;
    skr::Vector<ResourceNode*> resources;
    skr::Vector<BeforeExecuteCallback> exec_callbacks;

    friend struct TextureBuilder;
    friend struct BufferBuilder;
    friend struct AccelerationStructureBuilder;
    skr::ParallelFlatHashMap<CGPUBufferId, BufferHandle> imported_buffers;
    skr::ParallelFlatHashMap<CGPUTextureId, TextureHandle> imported_textures;
    skr::ParallelFlatHashMap<CGPUAccelerationStructureId, AccelerationStructureHandle> imported_acceleration_structures;
};
using RenderGraphSetupFunction = RenderGraph::RenderGraphSetupFunction;
using RenderGraphBuilder = RenderGraph::RenderGraphBuilder;
using RenderPassSetupFunction = RenderGraph::RenderPassSetupFunction;
using RenderPassBuilder = RenderGraph::RenderPassBuilder;
using ComputePassSetupFunction = RenderGraph::ComputePassSetupFunction;
using ComputePassBuilder = RenderGraph::ComputePassBuilder;
using CopyPassBuilder = RenderGraph::CopyPassBuilder;
using PresentPassSetupFunction = RenderGraph::PresentPassSetupFunction;
using PresentPassBuilder = RenderGraph::PresentPassBuilder;
using TextureSetupFunction = RenderGraph::TextureSetupFunction;
using TextureBuilder = RenderGraph::TextureBuilder;
using BufferSetupFunction = RenderGraph::BufferSetupFunction;
using BufferBuilder = RenderGraph::BufferBuilder;
using AccelerationStructureSetupFunction = RenderGraph::AccelerationStructureSetupFunction;
using AccelerationStructureBuilder = RenderGraph::AccelerationStructureBuilder;

struct PassInfoAnalysis;
struct QueueSchedule;
struct CrossQueueSyncAnalysis;
struct BarrierGenerationPhase;
struct MemoryAliasingPhase;
struct ResourceLifetimeAnalysis;
struct GraphViz
{
    static SKR_RENDER_GRAPH_API void generate_graphviz_visualization(
        skr::render_graph::RenderGraph* graph,
        const skr::render_graph::PassInfoAnalysis& info_analysis,
        const skr::render_graph::QueueSchedule& queue_schedule,
        const skr::render_graph::CrossQueueSyncAnalysis& ssis_phase,
        const skr::render_graph::BarrierGenerationPhase& barrier_phase,
        const skr::render_graph::MemoryAliasingPhase& aliasing_phase,
        const skr::render_graph::ResourceLifetimeAnalysis& lifetime_analysis
    );
};

} // namespace render_graph
} // namespace skr